// Copyright 2018 The Fuchsia Authors. All rights reserved.
// Use of this source code is governed by a BSD-style license that can be
// found in the LICENSE file.

#pragma once

#include <atomic>
#include <cstddef>

#include <lockdep/common.h>

namespace lockdep {

// A lock-free, wait-free hash set that tracks the set of lock classes that have
// been acquired prior the lock class that owns the set. Each lock class
// maintains its own dependency set.
//
// Implementation note: This hash set makes use of relaxed atomic operations.
// This approach is appropriate because the only variables communicated between
// threads are the values of the atomic variables themselves, other loads and
// stores are not published. Additionally, sequential consistency within a
// thread is ensured due to control dependencies only on the atomic variables.
class LockDependencySet {
 public:
  LockDependencySet() = default;
  LockDependencySet(const LockDependencySet&) = delete;
  void operator=(const LockDependencySet&) = delete;
  LockDependencySet(LockDependencySet&&) = delete;
  void operator=(LockDependencySet&&) = delete;

  // Checks the dependency hash set for the given lock class. This method may
  // safely race with AddLockClass(), converging on the correct answer by the
  // next check.
  bool HasLockClass(LockClassId id) const {
    for (size_t i = 0; i < kMaxLockDependencies; i++) {
      const auto& entry = GetEntry(id, i);
      const LockClassId entry_id = entry.load(std::memory_order_relaxed);
      if (entry_id == id)
        return true;
      else if (entry_id == kInvalidLockClassId)
        return false;
    }
    return false;
  }

  // Adds the given lock class id to the dependency set if not already present.
  // Updates the set using the following lock free approach:
  //   1. The dependency set is fixed size and all entries start out empty.
  //   2. New entries are added using open addressing with linear probing.
  //   3. An entry may only change from empty to holding a lock class id.
  //   4. To add an entry the set is probed linearly until either:
  //      a) The id to add appears in the set already.
  //      b) The first empty entry is found.
  //      c) The entire set has been probed; return max dependencies error.
  //   5. Attempt to compare-exchange the empty entry with the lock class id:
  //      a) If the update succeeds return success.
  //      b) If the update does not succeed but the winning value is the same
  //         lock class id return with success.
  //      c) If the update does not succeed and the winning value is a different
  //         lock class id proceed to the next entry and continue the probe from
  //         step #4.
  LockResult AddLockClass(LockClassId id) {
    for (size_t i = 0; i < kMaxLockDependencies; i++) {
      auto& entry = GetEntry(id, i);
      LockClassId entry_id = entry.load(std::memory_order_relaxed);
      if (entry_id == id)
        return LockResult::DependencyExists;
      while (entry_id == kInvalidLockClassId) {
        const bool result = entry.compare_exchange_weak(entry_id, id, std::memory_order_relaxed,
                                                        std::memory_order_relaxed);
        if (result) {
          return LockResult::Success;
        } else if (entry_id == id) {
          return LockResult::DependencyExists;
        } else {
          // Continue to find an unused slot, moving back to the outer
          // loop because of the while loop condition.
        }
      }
    }

    return LockResult::MaxLockDependencies;
  }

  // Iterator type to traverse the set of populated lock classes. Entries
  // added after the iterator is created may or may not be returned, depending
  // on where they land in the hash set relative to the index member.
  struct Iterator {
    const LockDependencySet* set;
    size_t index;

    LockClassId operator*() const {
      // TODO(eieio): See if it makes sense to add a memory order param
      // to the iterator.
      return set->list_[index].load(std::memory_order_relaxed);
    }

    Iterator operator++() {
      while (index != kMaxLockDependencies) {
        index++;
        if (**this != kInvalidLockClassId)
          break;
      }
      return *this;
    }

    bool operator!=(const Iterator& other) const {
      return set != other.set || index != other.index;
    }
  };

  // Iterator accessors.
  Iterator begin() const {
    Iterator iter{this, 0};
    if (*iter == kInvalidLockClassId)
      ++iter;
    return iter;
  }

  Iterator end() const { return {this, kMaxLockDependencies}; }

  // Clears the dependency set. This method is not used by the main algorithm
  // but may be useful for unit tests and benchmarking. Until lock sequence
  // memoization is implemented it is generally safe to call this method at
  // any time, resulting in dependency set being rebuilt at runtime. Once
  // lock sequence memoization is implemented it is necessary to clear the
  // memoization table whenever any dependency set is cleared so that the
  // dependency set can be rebuilt; failure to do so could result in missing
  // new lock violations.
  void clear() {
    for (size_t i = 0; i < kMaxLockDependencies; i++)
      list_[i].store(kInvalidLockClassId, std::memory_order_relaxed);
  }

 private:
  // Returns a reference to an entry by computing a trivial hash of the given id
  // and a linear probe offset.
  std::atomic<LockClassId>& GetEntry(LockClassId id, size_t offset) {
    const size_t index = (id + offset) % kMaxLockDependencies;
    return list_[index];
  }
  const std::atomic<LockClassId>& GetEntry(LockClassId id, size_t offset) const {
    const size_t index = (id + offset) % kMaxLockDependencies;
    return list_[index];
  }

  // The list of atomic variables that make up the hash set. Initialized to
  // kInvalidLockClassId  (0).
  std::atomic<LockClassId> list_[kMaxLockDependencies]{};
};

}  // namespace lockdep
